# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack
from datetime import timedelta
import functools
import hashlib
import json
import logging
import sys
from typing import Any
from typing import Dict
from typing import Optional
from typing import TextIO
from typing import Union

import anyio
from pydantic import BaseModel

from mcp import ClientSession
from mcp import StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client
from langchain_mcp_adapters.client import MultiServerMCPClient, load_mcp_tools

logger = logging.getLogger(__name__)


def retry_on_closed_resource(func):
    """Decorator to automatically retry action when MCP session is closed.

    When MCP session was closed, the decorator will automatically retry the
    action once. The create_session method will handle creating a new session
    if the old one was disconnected.

    Args:
        func: The function to decorate.

    Returns:
        The decorated function.
    """

    @functools.wraps(func)  # Preserves original function metadata
    async def wrapper(self, *args, **kwargs):
        try:
            return await func(self, *args, **kwargs)
        except anyio.ClosedResourceError:
            # Simply retry the function - create_session will handle
            # detecting and replacing disconnected sessions
            logger.info("Retrying %s due to closed resource", func.__name__)
            return await func(self, *args, **kwargs)

    return wrapper


class MCPSessionManager:
    """Manages MCP client sessions.

    This class provides methods for creating and initializing MCP client sessions,
    handling different connection parameters (Stdio and SSE) and supporting
    session pooling based on authentication headers.
    """

    def __init__(self):
        # Session pool: maps session keys to (session, exit_stack) tuples
        self._sessions: Dict[str, tuple[ClientSession, AsyncExitStack]] = {}
        # Lock to prevent race conditions in session creation
        self._session_lock = asyncio.Lock()

    def _is_session_disconnected(self, session: ClientSession) -> bool:
        """Checks if a session is disconnected or closed.

        Args:
            session: The ClientSession to check.

        Returns:
            True if the session is disconnected, False otherwise.
        """
        return session._read_stream._closed or session._write_stream._closed

    async def create_session(self, key: str, connection) -> ClientSession:
        # Use async lock to prevent race conditions
        async with self._session_lock:
            # Check if we have an existing session
            if key in self._sessions:
                session, exit_stack = self._sessions[key]

                # Check if the existing session is still connected
                if not self._is_session_disconnected(session):
                    # Session is still good, return it
                    return session
                else:
                    # Session is disconnected, clean it up
                    logger.info("Cleaning up disconnected session: %s", key)
                    try:
                        await exit_stack.aclose()
                    except Exception as e:
                        logger.warning(
                            "Error during disconnected session cleanup: %s", e
                        )
                    finally:
                        del self._sessions[key]

            # Create a new session (either first time or replacing disconnected one)
            exit_stack = AsyncExitStack()

            try:
                client = MultiServerMCPClient({key: connection})
                session = await exit_stack.enter_async_context(
                    client.session(key)
                )
                # Store session and exit stack in the pool
                self._sessions[key] = (session, exit_stack)
                logger.debug("Created new session: %s", key)
                return session

            except Exception:
                # If session creation fails, clean up the exit stack
                if exit_stack:
                    await exit_stack.aclose()
                raise

    async def close(self):
        """Closes all sessions and cleans up resources."""
        async with self._session_lock:
            for session_key in list(self._sessions.keys()):
                _, exit_stack = self._sessions[session_key]
                try:
                    await exit_stack.aclose()
                except Exception as e:
                    # Log the error but don't re-raise to avoid blocking shutdown
                    logger.error(
                        "Warning: Error during MCP session cleanup for"
                        f" {session_key}: {e}"
                    )
                finally:
                    del self._sessions[session_key]

    @retry_on_closed_resource
    async def get_tools(
        self,
        mcp_configs: dict
    ):
        """Return all tools in the toolset based on the provided context.

        Args:
            readonly_context: Context used to filter tools available to the agent.
                If None, all tools in the toolset are returned.

        Returns:
            List[BaseTool]: A list of tools available under the specified context.
        """
        # Get session from session manager
        tools = []
        for session_key, conn in mcp_configs.items():
            session = await self.create_session(session_key, conn)
            tools += await load_mcp_tools(session)

        return tools
